rm(list = ls())

## 

library(rdrobust)
library(tidyverse)
library(pbapply)
library(broom)
library(cowplot)
library(future.apply)
library(progressr)

## 

covs <- c('turnout_party_2009_btw', 'soz_vers_beschaeftigte_share',
          'pop_density_km2', 'migration_out_share')

## Get federal election data

bt <- readRDS("data/data_federal.rds") %>%
  filter(!is.na(treated)) 

## List of outcomes

outcomes <- c("turnout_party", 
              "agg_left_party", 
              "agg_center_party",
              "agg_right_party")

## Source helper function to tidy rdrobust output

source("code/tidy_rd.R")

## Drop 2009

bt <- bt %>% 
  filter(year > 2012)

## First differences

diff_df <- pblapply(outcomes, function(o) {
  out <- bt %>%
    filter(year > 2012) %>%
    pivot_wider(values_from = o, names_from = 'year', id_cols = 'ags',
                names_prefix = 'o') %>%
    mutate(diff = o2017  - o2013) %>%  dplyr::select(ags, diff) 
  ## Rename
  colnames(out)[2] <- o
  
  ## Return this
  out
}) %>%
  reduce(left_join) %>%
  left_join(bt %>% dplyr::select(ags, pop_dec_09, applies_census,
                                 one_of(covs), state_id) %>%
              distinct(ags, .keep_all = T)) %>%
  mutate(runvar = (pop_dec_09 * -1) + 10000) %>%
  filter(applies_census == 1)

## Placebo estimates

## First, set seed and declare placebo cutoffs

set.seed(1234)
n_cutoffs <- 1000
placebo_cutoff_list <- runif(n_cutoffs, -5000, 5000)

## Do in parallel via future.apply

# Define a progress handler
handlers(global = TRUE)
handlers("progress")

plan("multisession")

# Define the progressor function outside of the future_lapply function
with_progress({
  p <- progressor(steps = length(placebo_cutoff_list))
  out_rd_opt <- future_lapply(placebo_cutoff_list, function(cutoff) {
    cat(cutoff, '\n')
    p()  # Update the progress bar
    
    lapply(outcomes, function(o) {
      lapply(c(T,F), function(exclude_state) {
        if (exclude_state) {
          subset_select <- !diff_df$state_id == '08' 
        } else {
          subset_select <- rep(T, nrow(diff_df))
        }
        
        out <- rdrobust(y = diff_df[subset_select, o] %>% pull(!!o), 
                        x = diff_df$runvar[subset_select], c = cutoff,
                        covs = diff_df[subset_select, covs])
        
        ## Return this : Robust B-C SE
        out2 <- out %>% 
          tidy_rd(se_nr = 3) %>% 
          mutate(outcome = o, method = 'robust bias-corrected',
                 exclude_state = ifelse(exclude_state, 
                                        'State excluded', 'All states')) %>% 
          dplyr::rename(se = std.error, pval = p.value, 
                        bw = bw_left_h,
                        bw_bias = bw_left_b) %>% 
          mutate(conf.low90 = estimate - qnorm(0.95)*se,
                 conf.high90 = estimate + qnorm(0.95)*se)
        
        ## Both
        
        out2
      }) %>% reduce(rbind)
    }) %>% reduce(rbind)
  }, future.seed = 1234) %>% reduce(rbind)
})

## Distribution - BW excluded

results_main_exclude <- read_rds('data/rdd_federal_results.rds') %>% 
  filter(exclude_state == 'State excluded')

est_paper_left_exclude <- results_main_exclude %>% 
  filter(outcome == 'agg_left_party') %>% 
  pull(estimate)

est_paper_center_exclude <- results_main_exclude %>% 
  filter(outcome == 'agg_center_party') %>% 
  pull(estimate)

est_paper_right_exclude <- results_main_exclude %>% 
  filter(outcome == 'agg_right_party') %>% 
  pull(estimate)

est_paper_turnout_exclude <- results_main_exclude %>% 
  filter(outcome == 'turnout_party') %>% 
  pull(estimate)

## Placebo results with state excluded

out_bw_exclude <- out_rd_opt %>% 
  filter(exclude_state == 'State excluded')

# Figure 4: RI distribution ----

p_left <- ggplot(out_bw_exclude %>% 
                     filter(outcome == 'agg_left_party')) +
  geom_histogram(aes(x = estimate, y = after_stat(density)), 
                 bins =33,
                 fill = "grey97",
                 color = "black") + 
  geom_vline(xintercept = est_paper_left_exclude, linetype = 'dashed', size = 1) +
  scale_x_continuous(limits=c(-2, 2)) +
  xlab('RD estimate (p.p.)') + ylab('Density') +
  theme_bw() 
p_left

## Remaining outcomes
## Center

p_center <- ggplot(out_bw_exclude %>% filter(outcome == 'agg_center_party'), 
                   aes(x = estimate, y = after_stat(density))) +
  geom_histogram(fill = 'grey97', color = 'black', bins = 33) + 
  geom_vline(xintercept = est_paper_center_exclude, linetype = 'dashed', size = 1) +
  xlab('RD estimate (p.p.)') + ylab('Density') +
  theme_bw() 
p_center

## Right total

p_right <- ggplot(out_bw_exclude %>% filter(outcome == 'agg_right_party'), 
                  aes(x = estimate, y = after_stat(density))) +
  geom_histogram(fill = 'grey97', color = 'black', bins = 33) + 
  geom_vline(xintercept = est_paper_right_exclude, linetype = 'dashed', size = 1) +
  xlab('RD estimate (p.p.)') + ylab('Density') +
  theme_bw()
p_right

## Turnout

p_turn <- ggplot(out_bw_exclude %>% filter(outcome == 'turnout_party'), 
                 aes(x = estimate, y = after_stat(density))) +
  geom_histogram(fill = 'grey97', color = 'black', bins = 33) + 
  geom_vline(xintercept = est_paper_turnout_exclude, linetype = 'dashed', size = 1) +
  xlab('RD estimate (p.p.)') + ylab('Density') +
  theme_bw()
p_turn

# Figure A.11: distributions ----

labs <- c('Left parties,\nB-W excluded', 
          'Center parties,\nB-W excluded', 
          'Right parties,\nB-W excluded', 
          'Turnout,\nB-W excluded')

plot_grid(p_left, p_center, p_right,p_turn,  
          labels = LETTERS[1:4], ncol = 2)

## P-values

get_pval <- function(est_placebo, est_main) {
  if(est_main>0) {
    mean(abs(est_placebo) > est_main)
  } else {
    mean(-abs(est_placebo) < est_main)
  }
}

## Get p values for results that exclude BW

est_left_exclude <- out_bw_exclude %>% filter(outcome == 'agg_left_party') %>% 
  pull(estimate)
pval_left_exclude <- get_pval(est_left_exclude, est_paper_left_exclude)
pval_left_exclude

est_right_exclude <- out_bw_exclude %>% filter(outcome == 'agg_right_party') %>% 
  pull(estimate)
pval_right_exclude <-  get_pval(est_right_exclude, est_paper_right_exclude)
pval_right_exclude

est_center_exclude <- out_bw_exclude %>% filter(outcome == 'agg_center_party') %>% 
  pull(estimate)
pval_center_exclude <-  get_pval(est_center_exclude, est_paper_center_exclude)
pval_center_exclude

est_turnout_exclude <- out_bw_exclude %>% filter(outcome == 'turnout_party') %>% 
  pull(estimate)
pval_turnout_exclude <-  get_pval(est_turnout_exclude, est_paper_turnout_exclude)
pval_turnout_exclude

## Same for results based on not excluding B-W

results_main_include <- read_rds('data/rdd_federal_results.rds') %>% 
  filter(exclude_state != 'State excluded')

## Get estimates (including BW)

est_paper_left_include <- results_main_include %>% 
  filter(outcome == 'agg_left_party') %>% 
  pull(estimate)

est_paper_center_include <- results_main_include %>% 
  filter(outcome == 'agg_center_party') %>% 
  pull(estimate)

est_paper_right_include <- results_main_include %>% 
  filter(outcome == 'agg_right_party') %>% 
  pull(estimate)

est_paper_turnout_include <- results_main_include %>% 
  filter(outcome == 'turnout_party') %>% 
  pull(estimate)

## Get the placebo estimates

out_bw_include <- out_rd_opt %>% 
  filter(exclude_state != 'State excluded')

## Get p-values 

est_left_include <- out_bw_include %>% filter(outcome == 'agg_left_party') %>% 
  pull(estimate)
pval_left_include <- get_pval(est_left_include, est_paper_left_include)

est_center_include <- out_bw_include %>% filter(outcome == 'agg_center_party') %>% 
  pull(estimate)
pval_center_include <- get_pval(est_center_include, est_paper_center_include)

est_right_include <- out_bw_include %>% filter(outcome == 'agg_right_party') %>% 
  pull(estimate)
pval_right_include <- get_pval(est_right_include, est_paper_right_include)

est_turnout_include <- out_bw_include %>% filter(outcome == 'turnout_party') %>% 
  pull(estimate)
pval_turnout_include <- get_pval(est_turnout_include, est_paper_turnout_include)

## All pvals

pvals <- c(pval_left_exclude, 
           pval_center_exclude,
           pval_right_exclude, 
           pval_turnout_exclude,
           pval_left_include, 
           pval_center_include, 
           pval_right_include,
           pval_turnout_include)
exclude_state = c(rep('B-W excluded', 4),
                  rep('Full sample', 4))
outcome <- rep(c('Left parties', 'Center parties',
                 'Right parties', 'Turnout'), 2)

## Save pvals - they are used for another table that contains the main results

pvals <- data.frame(pval_ri = pvals, exclude_state,
                    outcome)
write_rds(pvals, 'data/ri_pvals.rds')
